import torch
import json
import codecs
import os
from contextlib import contextmanager
import torch.distributed as dist
import time
import numpy as np
import random
import pickle
from contextlib import contextmanager


def get_ms():
    return time.time() * 1000  # time.time()返回当前时间的时间戳（1970纪元后经过的浮点秒数）


def new_tensor(array, requires_grad=False):
    tensor = torch.tensor(array, requires_grad=requires_grad)  # 默认不允许优化
    if torch.cuda.is_available():
        tensor = tensor.cuda()  # 如果gpu能用，那么就放到GPU上
    return tensor


def init_seed(seed=None):
    if seed is None:
        seed = int(get_ms() // 1000)  # 如果不给定随机种子，那么用时间去生成

    np.random.seed(seed)
    random.seed(seed)

    torch.manual_seed(seed)  # Sets the seed for generating random numbers. Returns a torch.Generator object.
    torch.cuda.manual_seed(seed) # Sets the seed for generating random numbers for the current GPU. It’s safe to call this function if CUDA is not available; in that case, it is silently ignored.insufficient to get determinism
    torch.cuda.manual_seed_all(seed)  # 上面的方法不适用于多GPU；Sets the seed for generating random numbers on all GPUs.


def read_pkl(path):
    with codecs.open(path, 'rb') as f:
        return pickle.load(f)


def write_pkl(path, data):
    with codecs.open(path, 'wb') as f:
        pickle.dump(data, f)


def read_json(path):
    with codecs.open(path, encoding="utf-8") as f:
        return json.load(f)


def write_json(path, data):
    if not os.path.exists(os.path.dirname(path)):
        os.makedirs(os.path.dirname(path))
    with codecs.open(path, "w", encoding="utf-8") as f:
        json.dump(data, f)


def read_txt(path):
    with codecs.open(path, encoding="utf-8") as f:
        return [line[:-1] for line in f.readlines()]


def write_txt(path, data):
    if not os.path.exists(os.path.dirname(path)):
        os.makedirs(os.path.dirname(path))
    with codecs.open(path, "w", encoding="utf-8") as f:
        for line in data:
            f.write(str(line) + os.linesep)


def split_contres(datas):
    contexts, responses = [], []
    for data in datas:
        dia = data.split("\t")
        contexts.append(dia[:-1])
        responses.append(dia[-1])
    return contexts, responses

def random_index(k=4, pool_size=10, i0=0.5, i1=0.3, i2=0.2):
    # 以i0概率在0号位，以i1概率在1--(k-1)号位，以i2概率在k--(pool_size-1)号位
    total = 100
    i0, i1, i2 = i0*total, i1*total, i2*total
    t = random.randint(1, total)
    if k >= pool_size:
        if t <= 70:
            return 0
        else:
            return random.randint(1, pool_size-1) if pool_size > 1 else 0
    if t <= i0:
        return 0
    elif t <= i0+i1:
        return random.randint(1, k-1)
    else:
        return random.randint(k, pool_size-1)

def random_binary(i=0.7):
    # 以0.7概率生成1，以0.3概率生成0
    total = 100
    i = i*total
    t = random.randint(1, total)
    if t <= i:
        return 1
    else:
        return 0

if __name__ == "__main__":
    pass